# src/llm_client.py
import json
from typing import Dict, List, Any

from openai import OpenAI


class LLMClient:
    """
    负责与大语言模型 (LLM) API 进行通信，包括聊天和生成向量。
    """

    def __init__(self, base_url: str, api_key: str, chat_model: str, embedding_model: str):
        self.client = OpenAI(base_url=base_url, api_key=api_key)
        self.chat_model = chat_model
        self.embedding_model = embedding_model

    def _get_json_response(self, prompt: str) -> Dict[str, Any]:
        """
        向LLM发送请求并期望获得JSON格式的响应。
        """
        try:
            response = self.client.chat.completions.create(
                model=self.chat_model,
                messages=[{"role": "user", "content": prompt}],
                response_format={"type": "json_object"}
            )
            content = response.choices[0].message.content
            if content.strip().startswith("```json"):
                content = content.strip()[7:-3].strip()
            return json.loads(content)
        except Exception as e:
            print(f"Error calling LLM or parsing JSON: {e}")
            return {}

    def enrich_table_metadata(self, table_name: str, table_comment: str, columns: List[Dict[str, Any]]) -> Dict[
        str, str]:
        """
        为单个数据表生成语义化名称和描述，并将已有注释作为参考。
        """
        column_names = [col['column_name'] for col in columns]
        prompt = f"""
        You are a data cataloging expert. Your task is to provide semantic metadata for a database table.
        Analyze all the provided information to generate a comprehensive and accurate result.

        Input Information:
        - Table Name: '{table_name}'
        - Existing DBA Comment: '{table_comment or 'N/A'}'
        - Columns: {column_names}

        Task:
        Generate a JSON object with a 'semantic_name' and a 'description'.
        - 'semantic_name': A concise, human-readable name for the table in Chinese.
        - 'description': A detailed explanation in Chinese of the table's business purpose. If the existing comment is good, refine it; otherwise, create a new one based on the table and column names.

        Respond ONLY with a valid JSON object.
        """
        return self._get_json_response(prompt)

    def enrich_column_metadata(self, table_name: str, table_description: str, column: Dict[str, Any]) -> Dict[str, str]:
        """
        在表的上下文中，为单个列生成语义化名称和描述，并将已有注释作为参考。
        """
        prompt = f"""
        You are a data column expert. Given the context of a table, analyze the following database column.

        Context:
        - Table Name: '{table_name}'
        - Table Description: '{table_description}'
        
        Column to analyze:
        - Column Name: '{column['column_name']}'
        - Data Type: '{column['data_type']}'
        - Existing DBA Comment: '{column['column_comment'] or 'N/A'}'

        Task:
        Generate a JSON object with a 'semantic_name' and a 'description'.
        - 'semantic_name': A concise, human-readable name for the column in Chinese.
        - 'description': A detailed explanation in Chinese of this column's meaning. Refine the existing comment or create a new one based on all available information.
        
        Respond ONLY with a valid JSON object.
        """
        return self._get_json_response(prompt)

    def map_template_to_schema(self, template_headers: list, db_schema_string: str) -> dict:
        """
        让LLM分析模板表头和数据库Schema，生成字段映射关系。
        """
        prompt = f"""
        You are a data architect. Your task is to map the columns of a business report template to the physical columns in a database schema.

        **1. Database Schema:**
        (The following tables and columns are available.)
        ---
        {db_schema_string}
        ---

        **2. Desired Report Columns (Headers from the template file):**
        ---
        {template_headers}
        ---

        **Task:**
        For each header in the "Desired Report Columns", find the most appropriate database column (in `table_name.column_name` format) that would provide the data for it.
        - You must consider table joins to find the correct column.
        - If a report column like '销售总额' (Total Sales) clearly implies an aggregation, map it to the base column that would be aggregated (e.g., `transactions.price`). The aggregation logic itself will be handled later.
        - If a report column cannot be reasonably mapped to any database column, map its `db_column` to 'N/A'.

        Respond with a JSON object containing a single key "mappings". The value should be a list of objects, one for each header. Each object must have these keys:
        - "template_field": The exact header name from the report.
        - "db_column": The corresponding physical database column (`table.column`) or 'N/A'.
        - "description": A brief explanation in Chinese of why you chose this mapping.

        Example Response:
        {{
          "mappings": [
            {{ "template_field": "客户姓名", "db_column": "users.full_name", "description": "直接映射到用户表的全名" }},
            {{ "template_field": "销售总额", "db_column": "transactions.total_price", "description": "映射到交易表中的总价字段，后续可进行求和" }},
            {{ "template_field": "产品经理", "db_column": "N/A", "description": "数据库中没有找到与产品经理相关的信息" }}
          ]
        }}
        """
        response = self._get_json_response(prompt)
        return response.get("mappings", [])

    def adapt_sql_from_cache(self, user_query: str, cached_sql: str) -> str:
        """
        根据新的用户查询，微调一个缓存中已有的、相似的SQL。
        """
        prompt = f"""
        You are an expert SQL editor. Your task is to adapt a previously working SQL query to fit a new, slightly different user request.

        **1. Previously Working SQL:**
        (This query was successful for a similar request.)
        ---
        {cached_sql}
        ---

        **2. New User Request:**
        "{user_query}"

        **Task:**
        Modify the "Previously Working SQL" to accurately answer the "New User Request".
        Focus on changing specific values like names, dates, IDs, or numbers. Do not change the core logic (joins, aggregations) unless absolutely necessary.
        The resulting SQL must be complete and immediately executable.

        Respond with a JSON object containing a single key "sql".
        - "sql": A string containing the final, adapted SQL query.
        """
        response = self._get_json_response(prompt)
        return response.get("sql", "")

    def generate_sql_with_fixed_mappings(self, user_query: str, db_schema_string: str,
                                         mappings: List[Dict[str, Any]]) -> str:
        """
        在预设的字段映射约束下，根据用户查询生成最终SQL。
        """
        # 将映射关系格式化为字符串，方便LLM理解
        mappings_str = "\n".join([
            f"- Report Column '{m['template_field']}' MUST be generated from Database Column `{m['db_column']}`."
            for m in mappings
        ])

        prompt = f"""
        You are a SQL generation expert. Your task is to write a complete, executable SQL query based on a user's request, following a strict set of field mappings.

        **1. Database Schema:**
        (This is the database you can query from.)
        ---
        {db_schema_string}
        ---

        **2. Field Mappings (Constraints):**
        (You MUST follow these rules for the SELECT part of the query. Do not change them.)
        ---
        {mappings_str}
        ---

        **3. User's Specific Request:**
        "{user_query}"

        **Your Task:**
        Write a single, executable SQL query that answers the user's request.
        1.  The `SELECT` clause of your query is strictly defined by the field mappings. You should select the `db_column` values and alias them if necessary to match the spirit of the `template_field`.
        2.  Your main job is to figure out the necessary `FROM`, `JOIN`, `WHERE`, `GROUP BY`, and `ORDER BY` clauses to satisfy the user's request.
        3.  For example, if the user asks for 'Li Si' and a mapping exists for 'users.full_name', you must add a `WHERE users.full_name = 'Li Si'` clause.
        4.  If aggregations like 'Total Sales' are implied by a mapping to 'transactions.total_price', you must apply the correct aggregation (e.g., `SUM(transactions.total_price)`) and a corresponding `GROUP BY` clause.
        5.  If the filter condition is not clear, no need to add a `WHERE` clause.

        Respond with a JSON object containing one key: "sql".
        - "sql": A string containing the complete, final SQL query.
        """
        response = self._get_json_response(prompt)
        return response.get("sql", "")

    def generate_final_sql(self, db_schema_str: str, template_headers: list, matched_example_desc: str,
                           user_query: str) -> dict:
        """
        结合用户具体请求和匹配到的示例，生成最终可执行的SQL。
        """
        prompt = f"""
        You are an expert SQL developer. Your task is to write a final, executable SQL query based on a user's request, by adapting a known example use-case for a report.

        **1. Database Schema:**
        (The database you can query from)
        ---
        {db_schema_str}
        ---

        **2. Report Structure to Generate (Headers):**
        {template_headers}

        **3. Context from a Matched Example Use-Case:**
        You have found a template whose intended use is: "{matched_example_desc}"

        **4. User's Specific and New Request:**
        Now, generate a report for the following request: "{user_query}"

        **Task:**
        Analyze the user's new request and apply its specific entities (like names '李四', dates '九月份', regions, etc.) to the logic implied by the matched example. Write a single, complete, and executable SQL query.
        **DO NOT** use placeholders like `{{...}}`. The query must be ready to run directly on the database.
        The `SELECT` columns in your query must match the report headers.

        Respond with a JSON object containing one key: "sql".
        - "sql": A string containing the complete, final SQL query.
        """
        response = self._get_json_response(prompt)
        return response.get("sql", "")


    def get_embeddings(self, texts: List[str]) -> List[List[float]]:
        """
        使用 OpenAI Embedding API 为一批文本生成向量。
        """
        if not texts:
            return []

        try:
            # 替换空字符串，否则API会报错
            texts = [text.replace("\n", " ") for text in texts]
            response = self.client.embeddings.create(input=texts, model=self.embedding_model)
            return [embedding.embedding for embedding in response.data]
        except Exception as e:
            print(f"Error calling OpenAI Embedding API: {e}")
            return []
